"""
通用抽取智能体
"""

import json
from typing import Any, Dict

from . import DEFAULT_MODEL_NAME, client

USER_PROMPT = """对话记录：{text}
请你从上述对话记录中提取关键信息，以JSON格式返回。
"""


def extract_json_from_response(response: str) -> Dict[str, Any]:
    """
    解析LLM的JSON响应，提取关键参数

    Args:
        response: LLM的JSON格式响应文本

    Returns:
        UnemploymentParams对象，包含提取的参数
    """
    try:
        # 提取JSON字符串（可能在其他文本中）
        start_index = response.index("{")
        end_index = response.rindex("}") + 1

        data = json.loads(response[start_index:end_index])
        return data
    except (json.JSONDecodeError, AttributeError, ValueError):
        # 如果JSON解析失败，返回空对象
        return {}


async def run_agent(context: str, prompt_template: str) -> Dict[str, Any]:
    """
    运行参数提取agent

    Args:
        context: 输入的上下文文本

    Returns:
        JSON格式的参数提取结果
    """

    messages = [
        {"role": "system", "content": prompt_template},
        {"role": "user", "content": USER_PROMPT.format(text=context)},
    ]

    chat_completion = await client.chat.completions.create(
        messages=messages,
        model=DEFAULT_MODEL_NAME,
    )
    response_str = chat_completion.choices[0].message.content

    dic = extract_json_from_response(response_str)
    return dic
